{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7ae8d24d",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/vector_stores/Qdrant_hybrid_rag_multitenant_sharding.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ff04c3f",
   "metadata": {},
   "source": [
    "# Hybrid RAG with Qdrant: multi-tenancy, custom sharding, distributed setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40defd63",
   "metadata": {},
   "source": [
    "## What you'll build\n",
    "\n",
    "This notebook implements a production-style Hybrid RAG on Qdrant using LlamaIndex, designed for multitenancy and scale-out via custom sharding.\n",
    "\n",
    "- Hybrid search: dense embeddings + sparse BM25 for higher recall and precision.\n",
    "- Multitenancy: isolate tenants using payload filters and shard routing.\n",
    "- Custom sharding: keep each tenant local for performance and cost efficiency.\n",
    "- Distributed Qdrant: multi-node setup with replication for high availability and throughput."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d902b85b",
   "metadata": {},
   "source": [
    "This notebook walks through an end to end Retrieval Augmented Generation workflow that uses Qdrant as a distributed hybrid search backend and LlamaIndex as the orchestration layer. You will build a tenant aware RAG that combines dense vectors with sparse signals, you will isolate data per tenant with filters, and you will route data and queries with a custom shard key for scale."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3392404",
   "metadata": {},
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c2e7cab",
   "metadata": {},
   "source": [
    "### About the dependencies\n",
    "\n",
    "- llama-index: orchestration layer for ingestion, indexing, and retrieval.\n",
    "- llama-index-vector-stores-qdrant: Qdrant integration with hybrid support.\n",
    "- fastembed: lightweight CPU-friendly embedding/sparse models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87d3862c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U llama-index llama-index-vector-stores-qdrant fastembed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a580f4f",
   "metadata": {},
   "source": [
    "Make sure you have a distributed Qdrant cluster up and running. Here is a `compose.yaml` file:\n",
    "\n",
    "```yaml\n",
    "services:\n",
    "  qdrant_primary:\n",
    "    image: \"qdrant/qdrant:latest\"\n",
    "    ports:\n",
    "      - \"6333:6333\"\n",
    "    environment:\n",
    "      QDRANT__CLUSTER__ENABLED: \"true\"\n",
    "    command: [\"./qdrant\", \"--uri\", \"http://qdrant_primary:6335\"]\n",
    "    restart: always\n",
    "  qdrant_secondary:\n",
    "    image: \"qdrant/qdrant:latest\"\n",
    "    environment:\n",
    "      QDRANT__CLUSTER__ENABLED: \"true\"\n",
    "    command: [\"./qdrant\", \"--bootstrap\", \"http://qdrant_primary:6335\"]\n",
    "    restart: always\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79c07ade",
   "metadata": {},
   "source": [
    "## Imports and global settings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4feccad",
   "metadata": {},
   "source": [
    "### Settings and connectivity\n",
    "\n",
    "- Embeddings: `FastEmbedEmbedding('BAAI/bge-base-en-v1.5')` is a compact, high-quality baseline.\n",
    "- Connection: `QDRANT_URL` defaults to an HTTP endpoint; set `QDRANT_API_KEY` for secured/cloud setups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ad03414",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from qdrant_client import AsyncQdrantClient, QdrantClient\n",
    "from qdrant_client import models\n",
    "\n",
    "from llama_index.core import (\n",
    "    Settings,\n",
    "    VectorStoreIndex,\n",
    "    Document,\n",
    "    StorageContext,\n",
    ")\n",
    "from llama_index.vector_stores.qdrant import QdrantVectorStore\n",
    "from llama_index.embeddings.fastembed import FastEmbedEmbedding\n",
    "\n",
    "# Embeddings, small and fast\n",
    "Settings.embed_model = FastEmbedEmbedding(model_name=\"BAAI/bge-base-en-v1.5\")\n",
    "\n",
    "# Qdrant connection, local by default, set QDRANT_URL and QDRANT_API_KEY for cloud\n",
    "QDRANT_URL = os.getenv(\"QDRANT_URL\", \"http://localhost:6333\")\n",
    "QDRANT_API_KEY = os.getenv(\"QDRANT_API_KEY\")\n",
    "\n",
    "client: QdrantClient = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)\n",
    "aclient: AsyncQdrantClient = AsyncQdrantClient(\n",
    "    url=QDRANT_URL, api_key=QDRANT_API_KEY\n",
    ")\n",
    "COLLECTION = \"hybrid_rag_multitenant_sharding_demo\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebca4672",
   "metadata": {},
   "source": [
    "## Create distributed-ready collection"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "411e8e19",
   "metadata": {},
   "source": [
    "### Configure dual-vector schema (dense + sparse)\n",
    "\n",
    "- Define vector field names: `dense` for embeddings and `sparse` for BM25‑style signals.\n",
    "- Dense config:\n",
    "    - Determine embedding dimensionality at runtime by probing `Settings.embed_model` (avoids hardcoding).\n",
    "    - Use cosine distance for semantic similarity.\n",
    "- Sparse config:\n",
    "    - Enable an in‑memory sparse index (`on_disk=False`) to support hybrid scoring.\n",
    "- These settings establish the collection’s dual‑index layout used later by QdrantVectorStore for hybrid retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec6a935a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dense_vector_name = \"dense\"\n",
    "dense_config = models.VectorParams(\n",
    "    size=len(Settings.embed_model.get_text_embedding(\"probe\")),\n",
    "    distance=models.Distance.COSINE,\n",
    ")\n",
    "sparse_vector_name = \"sparse\"\n",
    "sparse_config = models.SparseVectorParams(\n",
    "    index=models.SparseIndexParams(on_disk=False)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da03d75b",
   "metadata": {},
   "source": [
    "### Shard keys and selector contract\n",
    "\n",
    "- `shard_keys`: ['tenant_a', 'tenant_b'] — predefined partitions used with custom sharding to keep each tenant local.\n",
    "- `payload_indexes`: keyword index on `tenant_id` to accelerate filter-based queries.\n",
    "- `shard_key_selector_fn(tenant_id) -> tenant_id`: returns the shard key used for both writes and reads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c1f0f06",
   "metadata": {},
   "outputs": [],
   "source": [
    "shard_keys = [\"tenant_a\", \"tenant_b\"]\n",
    "payload_indexes = [\n",
    "    {\n",
    "        \"field_name\": \"tenant_id\",\n",
    "        \"field_schema\": models.PayloadSchemaType.KEYWORD,\n",
    "    }\n",
    "]\n",
    "\n",
    "\n",
    "def shard_key_selector_fn(tenant_id: str) -> models.ShardKeySelector:\n",
    "    return tenant_id"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aea78d59",
   "metadata": {},
   "source": [
    "### Initialize hybrid Qdrant store with custom sharding\n",
    "\n",
    "This step creates or attaches to the collection named in `COLLECTION` and configures a dual‑vector hybrid store:\n",
    "\n",
    "- Hybrid search: `enable_hybrid=True` with `dense_vector_name='dense'` and `sparse_vector_name='sparse'`.\n",
    "- Dense config: `dense_config` uses cosine distance and derives size from `Settings.embed_model`.\n",
    "- Sparse config: `sparse_config` enables an in‑memory sparse index; `fastembed_sparse_model='Qdrant/bm25'` supplies BM25‑style signals.\n",
    "- Distributed topology:\n",
    "    - `sharding_method=Custom` with `shard_keys=['tenant_a','tenant_b']`.\n",
    "    - `shard_key_selector_fn(tenant_id) -> tenant_id` routes both writes and reads.\n",
    "    - `shard_number=6`, `replication_factor=2` for scale and High availability.\n",
    "- Payload index: `payload_indexes` accelerates filtering on `tenant_id`.\n",
    "\n",
    "Idempotent behavior: the vector store will create the collection if missing and reuse it on subsequent runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "926cc3b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store = QdrantVectorStore(\n",
    "    collection_name=COLLECTION,\n",
    "    client=client,\n",
    "    aclient=aclient,\n",
    "    dense_vector_name=dense_vector_name,\n",
    "    sparse_vector_name=sparse_vector_name,\n",
    "    enable_hybrid=True,\n",
    "    dense_config=dense_config,\n",
    "    sparse_config=sparse_config,\n",
    "    fastembed_sparse_model=\"Qdrant/bm25\",\n",
    "    shard_number=6,\n",
    "    sharding_method=models.ShardingMethod.CUSTOM,\n",
    "    shard_key_selector_fn=shard_key_selector_fn,\n",
    "    shard_keys=shard_keys,\n",
    "    replication_factor=2,\n",
    "    payload_indexes=payload_indexes,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "573a5cc2",
   "metadata": {},
   "source": [
    "## Prepare multi-tenant dataset\n",
    "We create two tenants with small document sets. Each Document carries tenant_id, tags, and a doc_id."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6b23a15",
   "metadata": {},
   "source": [
    "### Dataset design and extensibility\n",
    "\n",
    "We simulate two tenants with a few short documents each. Every `Document` carries:\n",
    "- `tenant_id` for isolation and shard routing,\n",
    "- `tags` for quick filtering and debugging,\n",
    "- `text` content used for dense/sparse indexing.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "678eac08",
   "metadata": {},
   "outputs": [],
   "source": [
    "TENANT_DOCS: dict[str, list[Document]] = {\n",
    "    \"tenant_a\": [\n",
    "        Document(\n",
    "            text=\"Solar panels reduce electricity bills and carbon footprint\",\n",
    "            metadata={\"tenant_id\": \"tenant_a\", \"tags\": [\"energy\", \"solar\"]},\n",
    "        ),\n",
    "        Document(\n",
    "            text=\"Inverters convert DC power to AC for home appliances\",\n",
    "            metadata={\"tenant_id\": \"tenant_a\", \"tags\": [\"energy\", \"hardware\"]},\n",
    "        ),\n",
    "        Document(\n",
    "            text=\"Net metering policies vary by region and utility provider\",\n",
    "            metadata={\n",
    "                \"tenant_id\": \"tenant_a\",\n",
    "                \"tags\": [\"policy\", \"regulation\"],\n",
    "            },\n",
    "        ),\n",
    "    ],\n",
    "    \"tenant_b\": [\n",
    "        Document(\n",
    "            text=\"Kubernetes orchestrates containers across a cluster\",\n",
    "            metadata={\"tenant_id\": \"tenant_b\", \"tags\": [\"cloud\", \"k8s\"]},\n",
    "        ),\n",
    "        Document(\n",
    "            text=\"Service meshes add observability and traffic management\",\n",
    "            metadata={\n",
    "                \"tenant_id\": \"tenant_b\",\n",
    "                \"tags\": [\"cloud\", \"networking\"],\n",
    "            },\n",
    "        ),\n",
    "        Document(\n",
    "            text=\"Helm charts package and deploy Kubernetes applications\",\n",
    "            metadata={\"tenant_id\": \"tenant_b\", \"tags\": [\"cloud\", \"devops\"]},\n",
    "        ),\n",
    "    ],\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e98d2f8",
   "metadata": {},
   "source": [
    "## Ingest with shard key for locality\n",
    "Here we embed text with the active Settings.embed_model, then upsert each point with payload and a shard key. This keeps each tenant local to a shard group in a cluster."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4734d2af",
   "metadata": {},
   "source": [
    "### Embedding strategy\n",
    "\n",
    "- FastEmbed keeps this demo CPU-friendly. For production, consider a service (e.g., text-embedding-3-large or in-house model) and cache embeddings.\n",
    "- If you change the model, update `dense_config.size` to match and consider reindexing.\n",
    "- Avoid embedding on every run in notebooks; persist or cache to speed up iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d8d4a86",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dense_embeddings(docs: list[Document]) -> list[Document]:\n",
    "    for doc in docs:\n",
    "        doc.embedding = Settings.embed_model.get_text_embedding(doc.text)\n",
    "    return docs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc68ad8e",
   "metadata": {},
   "source": [
    "### Ingestion flow and locality guarantees\n",
    "\n",
    "- We embed each doc with the configured embedding model (dense) and rely on the vector store to build the sparse representation.\n",
    "- Writes use `shard_identifier=tenant_id`, ensuring documents live on the intended shard group.\n",
    "\n",
    "Tip: For large batches, prefer the async ingestion APIs and chunk documents for backpressure control."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4086cdf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for tenant_id, docs in TENANT_DOCS.items():\n",
    "    docs = create_dense_embeddings(docs)\n",
    "    await vector_store.async_add(docs, shard_identifier=tenant_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce993d05",
   "metadata": {},
   "source": [
    "### Index wrapping and reusability\n",
    "\n",
    "`StorageContext.from_defaults(vector_store=vector_store)` binds the Qdrant collection to LlamaIndex's `VectorStoreIndex` without re-ingesting data.\n",
    "\n",
    "Benefits:\n",
    "- Reuse the same physical collection for multiple retrievers or query pipelines.\n",
    "- Swap retrieval modes (dense-only, sparse-only, hybrid) via retriever config, not data layout.\n",
    "- Keep ingestion concerns (sharding, replication) decoupled from application query logic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49ba6c04",
   "metadata": {},
   "outputs": [],
   "source": [
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "index = VectorStoreIndex.from_vector_store(\n",
    "    vector_store, storage_context=storage_context\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d645476",
   "metadata": {},
   "source": [
    "## Multi-tenant retrieval\n",
    "\n",
    "Use a tenant-scoped hybrid retriever and keep queries shard-local.\n",
    "You can also use metadata filters if you want to filter within the tenant's data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c1b2241",
   "metadata": {},
   "source": [
    "### Retrieval tips for hybrid mode\n",
    "\n",
    "- Set `vector_store_query_mode=HYBRID` to combine dense and sparse. Tune `similarity_top_k`, `sparse_top_k`, and `hybrid_top_k`.\n",
    "- Pass `vector_store_kwargs={\"shard_identifier\": tenant_id}` to keep queries within the tenant's shard.\n",
    "- Add metadata filters (e.g., on `tenant_id` or `tags`) to further narrow candidates when needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5131fc1b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tenant: tenant_b | Query: manage microservices traffic and observability\n",
      "1. score=4.6271 | tags=['cloud', 'networking'] | text=Service meshes add observability and traffic management\n",
      "2. score=0.1213 | tags=['cloud', 'k8s'] | text=Kubernetes orchestrates containers across a cluster\n",
      "3. score=0.0000 | tags=['cloud', 'devops'] | text=Helm charts package and deploy Kubernetes applications\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.retrievers import VectorIndexRetriever\n",
    "from llama_index.core.vector_stores.types import VectorStoreQueryMode\n",
    "\n",
    "\n",
    "def create_retriever_for_tenant(tenant_id: str) -> VectorIndexRetriever:\n",
    "    if tenant_id not in shard_keys:\n",
    "        raise ValueError(\n",
    "            f\"Unknown tenant_id: {tenant_id}. Expected one of {shard_keys}\"\n",
    "        )\n",
    "    return VectorIndexRetriever(\n",
    "        index=index,\n",
    "        vector_store_query_mode=VectorStoreQueryMode.HYBRID,\n",
    "        similarity_top_k=5,\n",
    "        sparse_top_k=5,\n",
    "        hybrid_top_k=5,\n",
    "        vector_store_kwargs={\"shard_identifier\": tenant_id},\n",
    "    )\n",
    "\n",
    "\n",
    "tenant_id = \"tenant_b\"\n",
    "retriever = create_retriever_for_tenant(tenant_id)\n",
    "\n",
    "query = \"manage microservices traffic and observability\"\n",
    "results = retriever.retrieve(query)\n",
    "\n",
    "print(f\"Tenant: {tenant_id} | Query: {query}\")\n",
    "for i, r in enumerate(results, 1):\n",
    "    meta = r.node.metadata\n",
    "    print(\n",
    "        f\"{i}. score={r.score:.4f} | tags={meta.get('tags')} | text={r.node.get_content()}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72770851",
   "metadata": {},
   "source": [
    "### Interpreting results\n",
    "\n",
    "- The printout shows the hybrid score, tags (metadata), and snippet of the matched text.\n",
    "- Verify tenant isolation by switching `tenant_id` and observing that results come only from that tenant's documents."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "edap",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
